#include "flux.h"

using namespace std;
using namespace mu;

//------------------------------------------------------------------------------
Flux::Flux(Mesh* mymesh,FsContainer* myfs,ClassParameter& param){

  nbProcs = param.getInteger("NbThreads");
  mesh    = mymesh;
  fs      = myfs;
  nbEq    = 1;                                                                  //NOT YET IMPLEMENTED for equation >1
    
  getBondaryConditions(param);
  
}
//------------------------------------------------------------------------------
void Flux::getBondaryConditions(ClassParameter& param){
  
  int nbFNP = (*fs)[0]->getNbFNP();
  int nbB   = mesh->getNbBFaces();
  
  UBNP = DGMatrix(nbFNP,nbB);
  
  string bcName = param.getToken("BConditions");
  if(bcName=="Dirichlet") BCtype=Dirichlet;
  if(bcName=="Periodic")  BCtype=Periodic;

  if(BCtype==Dirichlet){
    
    vector<int>    BCgmshTags = param.getIntegerList("BCTags");
    vector<string> BCfunction = param.getStringList("BCValues");
    map<int,MathParser> expression;
    
    double x=0.;
    double y=0.;
    
    int bT=0;
    for(vector<int>::const_iterator BCTagIT=BCgmshTags.begin();BCTagIT!=BCgmshTags.end();++BCTagIT,bT++){
      MathParser tmp; 
      expression[*BCTagIT] = tmp; 
      expression[*BCTagIT].DefineVar("x",&x); 
      expression[*BCTagIT].DefineVar("y",&y); 
      expression[*BCTagIT].SetExpr(BCfunction[bT]);
    }
    
    const DGIntMatrix* FRefNodesIndex = (*fs)[0]->getFRefNodesIndex();
    
    for(int bID=0;bID<nbB;bID++){
      int BCTag  = mesh->getBFacesTag(bID,3);
      int eTL    = mesh->getBFacesTag(bID,0);
      int eIDL   = mesh->getBFacesTag(bID,1);
      int eFnumL = mesh->getBFacesTag(bID,2);
      int Fnode0 = (*FRefNodesIndex)(eFnumL,0);
      int Fnode1 = (*FRefNodesIndex)(eFnumL,nbFNP-1);
      const DGMatrix* NPcoord = mesh->getNPcoord(eTL,eIDL);
      double Vx  = (*NPcoord)(Fnode1,0)-(*NPcoord)(Fnode0,0);
      double Vy  = (*NPcoord)(Fnode1,1)-(*NPcoord)(Fnode0,1);
      for(int n=0;n<nbFNP;n++){
        x = (*NPcoord)((*FRefNodesIndex)(eFnumL,n),0);
        y = (*NPcoord)((*FRefNodesIndex)(eFnumL,n),1);
        UBNP(n,bID) = expression[BCTag].Eval();
      }
    }
    
  }
  
  
}
//------------------------------------------------------------------------------
LinearTransport::LinearTransport(Mesh* mymesh,FsContainer* myfs,ClassParameter& param):Flux(mymesh,myfs,param){
  
  // --- build Beta ---
  
  BetaMax = 0.;
  BetaMin = 10000.;
  maxBetaOverHmax = 0.;
  maxBetaOverHmin = 0.;
  
  vector<string> BetaStr = param.getStringList("Beta");

  double BetaNorm = 0.;
  double x        = 0.;
  double y        = 0.;
  BetaMax         = 0.;

  MathParser expressionX;
  MathParser expressionY;
  expressionX.DefineVar("x", &x);
  expressionY.DefineVar("x", &x); 
  expressionX.DefineVar("y", &y); 
  expressionY.DefineVar("y", &y); 
  expressionX.SetExpr(BetaStr[0]);
  expressionY.SetExpr(BetaStr[1]);
    
  int eT=0;
  for(FsContainer::const_iterator fsIt = fs->begin();fsIt!=fs->end();++fsIt,eT++){
    
    const DGMatrix*    ERefNodesPos   = (*fsIt)->getERefNodesPos();
    const DGIntMatrix* FRefNodesIndex = (*fsIt)->getFRefNodesIndex();
    
    const DGMatrix* EQPPos       = (*fsIt)->getEQP();
    const DGVector* FQPPos       = (*fsIt)->getFQP();
    FunctionSpace* mapping       = FunctionSpaceFactory::CreateNewFunctionSpace((FctType)mesh->getEGmshType(eT),1);
    int nbENP                    = (*fsIt)->getNbENP();
    int nbFNP                    = (*fsIt)->getNbFNP();
    int nbEQP                    = (*fsIt)->getNbEQP();
    int nbFQP                    = (*fsIt)->getNbFQP();
    int nbF                      = mesh->getNbElmtFaces(eT);
    int nbElmt                   = mesh->getNbElmt(eT);
    int nbFace                   = mesh->getNbIntFaces();
    int nbBFace                  = mesh->getNbBFaces();
    
    BetaENP.push_back(DGMatrix(nbENP,nbElmt*2));
    BetaEQP.push_back(DGMatrix(nbEQP,nbElmt*2));
    BetaFNP.push_back(DGMatrix(nbFNP,nbFace*2));
    BetaFQP.push_back(DGMatrix(nbFQP,nbFace*2));
    BetaBNP.push_back(DGMatrix(nbFNP,nbBFace*2));
    BetaBQP.push_back(DGMatrix(nbFQP,nbBFace*2));
    
    // --- loop over the element ---
    
    //#pragma omp parallel for
    for(int eID=0;eID<nbElmt;eID++){
      double hmin = mesh->getElmtH(eT,eID,0);
      double hmax = mesh->getElmtH(eT,eID,1);
      for(int n=0;n<nbENP;n++){
        x = 0.; y = 0.;
        for(int j=0;j<nbF;j++){
          double phi_j = mapping->getPhiE(j,(*ERefNodesPos)(n,0),(*ERefNodesPos)(n,1)); 
          x += mesh->getNode(mesh->getElmtNodesTag(eT,eID,j),0)*phi_j;
          y += mesh->getNode(mesh->getElmtNodesTag(eT,eID,j),1)*phi_j;
        }
        BetaENP[eT](n,eID*2+0) = expressionX.Eval();
        BetaENP[eT](n,eID*2+1) = expressionY.Eval();
      }  
      for(int q=0;q<nbEQP;q++){
        x = 0.; y = 0.;
        for(int j=0;j<nbF;j++){
          double phi_j = mapping->getPhiE(j,(*EQPPos)(q,0),(*EQPPos)(q,1)); 
          x += mesh->getNode(mesh->getElmtNodesTag(eT,eID,j),0)*phi_j;
          y += mesh->getNode(mesh->getElmtNodesTag(eT,eID,j),1)*phi_j;
        }
        BetaEQP[eT](q,eID*2+0) = expressionX.Eval();
        BetaEQP[eT](q,eID*2+1) = expressionY.Eval();
        double BetaNorm = sqrt(BetaEQP[eT](q,eID*2+0)*BetaEQP[eT](q,eID*2+0)+BetaEQP[eT](q,eID*2+1)*BetaEQP[eT](q,eID*2+1));
        if(BetaMax<BetaNorm) BetaMax = BetaNorm;
        if(BetaMin>BetaNorm) BetaMin = BetaNorm;
        if(maxBetaOverHmax<BetaNorm/hmax) maxBetaOverHmax=BetaNorm/hmax;
        if(maxBetaOverHmin<BetaNorm/hmin) maxBetaOverHmin=BetaNorm/hmin;        
      } 
    }
    
    // --- loop over the interfaces ---
    
    //#pragma omp parallel for
    for(int fID=0;fID<nbFace;fID++){
      int eTL    = mesh->getIntFacesTag(fID,0);
      int eIDL   = mesh->getIntFacesTag(fID,1);
      int eFnumL = mesh->getIntFacesTag(fID,2);
      int Fnode0 = (*FRefNodesIndex)(eFnumL,0);
      int Fnode1 = (*FRefNodesIndex)(eFnumL,nbFNP-1);
      const DGMatrix* NPcoord = mesh->getNPcoord(eTL,eIDL);
      double Vx  = (*NPcoord)(Fnode1,0)-(*NPcoord)(Fnode0,0);
      double Vy  = (*NPcoord)(Fnode1,1)-(*NPcoord)(Fnode0,1);
      double hmin = mesh->getElmtH(eT,eIDL,0);
      double hmax = mesh->getElmtH(eT,eIDL,1);
      for(int n=0;n<nbFNP;n++){
        x = (*NPcoord)((*FRefNodesIndex)(eFnumL,n),0);
        y = (*NPcoord)((*FRefNodesIndex)(eFnumL,n),1);
        BetaFNP[eT](n,fID*2+0) = expressionX.Eval();
        BetaFNP[eT](n,fID*2+1) = expressionY.Eval();
      }
      for(int q=0;q<nbFQP;q++){
        x = (*NPcoord)(Fnode0,0) + 0.5*((*FQPPos)(q)+1)*Vx;                     // 0.5*(+1) because FQP are in [-1,1] and we want to be in [0,1]
        y = (*NPcoord)(Fnode0,1) + 0.5*((*FQPPos)(q)+1)*Vy;
        BetaFQP[eT](q,fID*2+0) = expressionX.Eval();
        BetaFQP[eT](q,fID*2+1) = expressionY.Eval();
        double BetaNorm = sqrt(BetaFQP[eT](q,fID*2+0)*BetaFQP[eT](q,fID*2+0)+BetaFQP[eT](q,fID*2+1)*BetaFQP[eT](q,fID*2+1));
        if(BetaMax<BetaNorm) BetaMax = BetaNorm;
        if(BetaMin>BetaNorm) BetaMin = BetaNorm;
        if(maxBetaOverHmax<BetaNorm/hmax) maxBetaOverHmax=BetaNorm/hmax;
        if(maxBetaOverHmin<BetaNorm/hmin) maxBetaOverHmin=BetaNorm/hmin; 
      }
    }
    
    // --- loop over the boundary ---
    
    //#pragma omp parallel for
    for(int bID=0;bID<nbBFace;bID++){
      int eTL    = mesh->getBFacesTag(bID,0);
      int eIDL   = mesh->getBFacesTag(bID,1);
      int eFnumL = mesh->getBFacesTag(bID,2);
      int Fnode0 = (*FRefNodesIndex)(eFnumL,0);
      int Fnode1 = (*FRefNodesIndex)(eFnumL,nbFNP-1);
      const DGMatrix* NPcoord = mesh->getNPcoord(eTL,eIDL);
      double Vx  = (*NPcoord)(Fnode1,0)-(*NPcoord)(Fnode0,0);
      double Vy  = (*NPcoord)(Fnode1,1)-(*NPcoord)(Fnode0,1);
      double hmin = mesh->getElmtH(eT,eIDL,0);
      double hmax = mesh->getElmtH(eT,eIDL,1);
      for(int n=0;n<nbFNP;n++){
        x = (*NPcoord)((*FRefNodesIndex)(eFnumL,n),0);
        y = (*NPcoord)((*FRefNodesIndex)(eFnumL,n),1);
        BetaBNP[eT](n,bID*2+0) = expressionX.Eval();
        BetaBNP[eT](n,bID*2+1) = expressionY.Eval();
      }
      for(int q=0;q<nbFQP;q++){
        x = (*NPcoord)(Fnode0,0) + 0.5*((*FQPPos)(q)+1)*Vx;                     
        y = (*NPcoord)(Fnode0,1) + 0.5*((*FQPPos)(q)+1)*Vy;
        BetaBQP[eT](q,bID*2+0) = expressionX.Eval();
        BetaBQP[eT](q,bID*2+1) = expressionY.Eval();
        double BetaNorm = sqrt(BetaBQP[eT](q,bID*2+0)*BetaBQP[eT](q,bID*2+0)+BetaBQP[eT](q,bID*2+1)*BetaBQP[eT](q,bID*2+1));
        if(BetaMax<BetaNorm) BetaMax = BetaNorm;
        if(BetaMin>BetaNorm) BetaMin = BetaNorm;
        if(maxBetaOverHmax<BetaNorm/hmax) maxBetaOverHmax=BetaNorm/hmax;
        if(maxBetaOverHmin<BetaNorm/hmin) maxBetaOverHmin=BetaNorm/hmin; 
      }
    }
    
    // ---
    
    delete mapping;
    
  }
  
};
//------------------------------------------------------------------------------
void LinearTransport::getFlux(int dim,const DGMatrix* U,DGMatrix* flux,PType pType,int eT,int IDimposed){
  
  SolContainer* Beta;
  switch (pType) {
    case EQP: Beta = &BetaEQP; break;                                           
    case FQP: Beta = &BetaFQP; break;
    case BQP: Beta = &BetaBQP; break;
    case ENP: Beta = &BetaENP; break;
    case FNP: Beta = &BetaFNP; break;
    case BNP: Beta = &BetaBNP; break;
  }

  if(U->size2()==1){
    for(int j=0;j<U->size2();j++){
      for(int i=0;i<U->size1();i++){
        (*flux)(i,j) = (*Beta)[eT](i,IDimposed*2+dim)*(*U)(i,j);
      }
    }
  }else{
    #pragma omp parallel for
    for(int j=0;j<U->size2();j++){
      for(int i=0;i<U->size1();i++){
        (*flux)(i,j) = (*Beta)[eT](i,j*2+dim)*(*U)(i,j);
      }
    }
  }

}
//------------------------------------------------------------------------------
void LinearTransport::getFluxDerivative(int dim,const DGMatrix* U,DGMatrix* flux,PType pType,int eT,int IDimposed){
  
  SolContainer* Beta;
  switch (pType) {
    case EQP: Beta = &BetaEQP; break;                                           
    case FQP: Beta = &BetaFQP; break;
    case BQP: Beta = &BetaBQP; break;
    case ENP: Beta = &BetaENP; break;
    case FNP: Beta = &BetaFNP; break;
    case BNP: Beta = &BetaBNP; break;
  }
  
  if(U->size2()==1){
    for(int j=0;j<U->size2();j++){
      for(int i=0;i<U->size1();i++){
        (*flux)(i,j) = (*Beta)[eT](i,IDimposed*2+dim);
      }
    }
  }else{
    #pragma omp parallel for
    for(int j=0;j<U->size2();j++){
      for(int i=0;i<U->size1();i++){
        (*flux)(i,j) = (*Beta)[eT](i,j*2+dim);
      }
    }
  }
  
}
//------------------------------------------------------------------------------
void UpwindLinearTransport::getNumericalFlux(const DGMatrix* UL,const SolContainer* FUL,const SolContainer* dFUL,
                                             const DGMatrix* UR,const SolContainer* FUR,const SolContainer* dFUR,
                                             DGMatrix* flux,const DGVector* normal,int fID){
  
  DGMatrix jump(UL->size1(),UL->size2());
  jump.memcpy(*UL);
  jump -= (*UR);
  
  double FluxPrim = 0;
  for(int i=0;i<UL->size1();i++){
    FluxPrim = max(FluxPrim,fabs((*dFUL)[0](i,fID)*(*normal)(0) + (*dFUL)[1](i,fID)*(*normal)(1)));
  }
  for(int i=0;i<UL->size1();i++){
    //double FluxPrim = fabs((*dFUL)[0](i,fID)*(*normal)(0) + (*dFUL)[1](i,fID)*(*normal)(1));
    (*flux)(i,0)    = 0.5*(((*FUL)[0](i,fID)+(*FUR)[0](i,fID))*(*normal)(0) + ((*FUL)[1](i,fID)+(*FUR)[1](i,fID))*(*normal)(1))
                    + 0.5*FluxPrim*jump(i,0);
  }
                                                                              
}
//------------------------------------------------------------------------------
void CenterLinearTransport::getNumericalFlux(const DGMatrix* UL,const SolContainer* FUL,const SolContainer* dFUL,
                                             const DGMatrix* UR,const SolContainer* FUR,const SolContainer* dFUR,
                                             DGMatrix* flux,const DGVector* normal,int fID){
    
  //Centered
  for(int i=0;i<UL->size1();i++){
    (*flux)(i,0) = 0.5*(((*FUL)[0](i,fID)+(*FUR)[0](i,fID))*(*normal)(0) + ((*FUL)[1](i,fID)+(*FUR)[1](i,fID))*(*normal)(1));
  }
  
}
//------------------------------------------------------------------------------
KPP::KPP(Mesh* mymesh,FsContainer* myfs,ClassParameter& param):Flux(mymesh,myfs,param){}
//------------------------------------------------------------------------------
void KPP::getFlux(int dim,const DGMatrix* U,DGMatrix* flux,PType pType,int eT,int IDimposed){
  
  for(int j=0;j<U->size2();j++){
    for(int i=0;i<U->size1();i++){
      switch (dim) {
        case 0: (*flux)(i,j) = sin((*U)(i,j)); break;
        case 1: (*flux)(i,j) = cos((*U)(i,j)); break;
      }
    }
  }
  
}
//------------------------------------------------------------------------------
void KPP::getFluxDerivative(int dim,const DGMatrix* U,DGMatrix* flux,PType pType,int eT,int IDimposed){
  
  for(int j=0;j<U->size2();j++){
    for(int i=0;i<U->size1();i++){
      switch (dim) {
        case 0: (*flux)(i,j) =  cos((*U)(i,j)); break;
        case 1: (*flux)(i,j) = -sin((*U)(i,j)); break;
      }
    }
  }
  
}
//------------------------------------------------------------------------------
void UpwindKPP::getNumericalFlux(const DGMatrix* UL,const SolContainer* FUL,const SolContainer* dFUL,
                                             const DGMatrix* UR,const SolContainer* FUR,const SolContainer* dFUR,
                                             DGMatrix* flux,const DGVector* normal,int fID){
  
  DGMatrix jump(UL->size1(),UL->size2());
  jump.memcpy(*UL);
  jump -= (*UR);
  
  //Upwind
  double FluxPrim = 0;
  for(int i=0;i<UL->size1();i++){
    FluxPrim = max(FluxPrim,fabs((*dFUL)[0](i,fID)*(*normal)(0) + (*dFUL)[1](i,fID)*(*normal)(1)));
  }
  for(int i=0;i<UL->size1();i++){
    (*flux)(i,0)    = 0.5*(((*FUL)[0](i,fID)+(*FUR)[0](i,fID))*(*normal)(0) + ((*FUL)[1](i,fID)+(*FUR)[1](i,fID))*(*normal)(1))
    + 0.5*FluxPrim*jump(i,0);
  }
  
}
//------------------------------------------------------------------------------
void CenterKPP::getNumericalFlux(const DGMatrix* UL,const SolContainer* FUL,const SolContainer* dFUL,
                                             const DGMatrix* UR,const SolContainer* FUR,const SolContainer* dFUR,
                                             DGMatrix* flux,const DGVector* normal,int fID){
  
  //Centered
  for(int i=0;i<UL->size1();i++){
    (*flux)(i,0) = 0.5*(((*FUL)[0](i,fID)+(*FUR)[0](i,fID))*(*normal)(0) + ((*FUL)[1](i,fID)+(*FUR)[1](i,fID))*(*normal)(1));
  }
  
}
//------------------------------------------------------------------------------
Flux* FluxFactory::CreateNewFlux(Mesh* mymesh,FsContainer* myfs,ClassParameter& param) {
  
  string pde           = param.getToken("PDE");
  string graphL        = param.getToken("GraphLaplacian");
  string lumpedM       = param.getToken("LumpedMass");
  string numericalflux = param.getToken("NumericalFlux");

  cout << "[PDE] - flux definition --------> " << pde << endl;	
  cout << "[PDE] - numerical flux  --------> " << numericalflux << endl;	
  cout << "[PDE] - graph Laplacian --------> " << graphL << endl;	
  cout << "[PDE] - lumped Mass ------------> " << lumpedM << endl;	
  
  if(pde=="LinearTransport" && numericalflux=="Upwind") return new UpwindLinearTransport(mymesh,myfs,param);
  if(pde=="LinearTransport" && numericalflux=="Center") return new CenterLinearTransport(mymesh,myfs,param);
  if(pde=="KPP"             && numericalflux=="Upwind") return new UpwindKPP(mymesh,myfs,param);
  if(pde=="KPP"             && numericalflux=="Center") return new CenterKPP(mymesh,myfs,param);  
  
}